# -*- coding: utf-8 -*-

"""
    将原始数据集进行划分成训练集、验证集和测试集
"""

import random
import pandas as pd
import numpy as np
from utils.utils import get_classes

#--------------------------------------------#
#   train_per, valid_per, test_per分别用于
#   指定数据集中训练集, 验证集和测试集的比例
#   默认情况下比例为8:1:1
#--------------------------------------------#
train_per    = 0.8
valid_per    = 0.2
test_per     = 0.0

classes_path    = 'model_data/news_classes.txt'

classes, num_classes = get_classes(classes_path)

#-------------------------------------------------------#
#   datasets_path   指向数据集所在的文件
#-------------------------------------------------------#
dataset_path = "image_true_label_train.csv"

train_csv_path = "train_Dataset.csv"
valid_csv_path = "valid_Dataset.csv"
test_csv_path  = "test_Dataset.csv"



if __name__ == '__main__':

    #--------------------------------------------#
    #   划分训练集, 验证集和测试集
    #   生成train, valid, test三个文件
    #--------------------------------------------#
    print("Generate train, valid, test csv.")

    # 读取csv文件
    CsvData = pd.read_csv(dataset_path)
    dataCsvColumns = CsvData.columns[1:]
    Data = np.array(CsvData.iloc[:, 1:])

    # 生成train, valid, test三个列表
    for i in range(num_classes):
        idata = Data[Data[:, -1] == i]
        random.seed(666)
        random.shuffle(idata)
        idata_num = len(idata)
        train_point = int(idata_num * train_per)
        valid_point = int(idata_num * (train_per + valid_per))
        if i == 0:
            train_csv = idata[:train_point, :]
            valid_csv = idata[train_point:valid_point, :]
            test_csv = idata[valid_point:, :]
        else:
            train_csv = np.concatenate((train_csv, idata[:train_point, :]), axis=0)
            valid_csv = np.concatenate((valid_csv, idata[train_point:valid_point, :]), axis=0)
            test_csv = np.concatenate((test_csv, idata[valid_point:, :]), axis=0)

    #---------------------------------- #
    #   生成train, valid, test三个csv文件
    #---------------------------------- #
    print("Generate csv.")
    train_data = pd.DataFrame(train_csv, columns=dataCsvColumns)
    train_data.to_csv(train_csv_path, index=False)
    valid_data = pd.DataFrame(valid_csv, columns=dataCsvColumns)
    valid_data.to_csv(valid_csv_path, index=False)
    test_data = pd.DataFrame(test_csv, columns=dataCsvColumns)
    test_data.to_csv(test_csv_path, index=False)
    print("Generate csv done.")


